#include "networkfactory.h"

#include "nodes/normalnode.h"
#include "nodes/gammanode.h"
#include "nodes/gammapinode.h"
#include "nodes/inversegammanode.h"
#include "nodes/betanode.h"
#include "nodes/bernoullinode.h"
#include "nodes/bernoullidirectnode.h"
#include "nodes/binomialnode.h"
#include "nodes/poissonnode.h"
#include "nodes/combinationnode.h"

#include <QDebug>
#include <QFile>
#include <QStringList>
#include <QMap>

#include <assert.h>
#include <qmath.h>
#include <gsl/gsl_statistics.h>


Network * NetworkFactory::buildNetwork(NetworkFactory::NetworkType type)
{
    switch(type)
    {
        case FACULTY:
            return buildFacultyNetwork();

        case BURGLAR:
            return buildBurglarNetwork();

        case SIMPLE_BERNOULLI:
            return buildSimpleBernoulli();

        case GOLFERS:
            return buildGolfersNetwork();

        case WACKY:
            return buildWackyNetwork();

        case POPCORN:
            return buildPopcornNetwork();

        default:
            assert(false);
    }
}


/******************************************************************************************************************************************************
 *
 *   Faculty Evaluations
 *
 *   In an example in class, we modeled faculty evaluation scores as distributed normally with a normal prior over the mean parameter
 *   and an inverse gamma prior over the variance parameter. Therefore there are two parameter nodes and N observed nodes (one for each datum).
 *   Andrew's prior belief is represented by a Normal(5, 1/9) prior distribution for mean and an InvGamma(mean=1/4, var=1/144) for variance
 *   (alpha=11, beta=2.5). Each piece of data is distributed according to N(mean, var). Since you're implementing Gamma, note that an inverse
 *   gamma random sample is 1 divided by a gamma random sample.
 *
 ********************************************************************************************************************************************************/

Network * NetworkFactory::buildFacultyNetwork()
{
    Network * network = new Network();

    double numPoints = 23;
    double data[] = {
                        6.39,
                        6.32,
                        6.25,
                        6.24,
                        6.21,
                        6.18,
                        6.17,
                        6.13,
                        6.00,
                        6.00,
                        5.97,
                        5.82,
                        5.81,
                        5.71,
                        5.55,
                        5.50,
                        5.39,
                        5.37,
                        5.35,
                        5.30,
                        5.27,
                        4.94,
                        4.50
                    };

    double sigma    = sqrt(gsl_stats_variance(data, 1, 23));
    double mean     = gsl_stats_mean(data, 1, 23);

    // mu node
    NormalNode * mu = new NormalNode("mu", mean, sigma,
                               new Node(5.0),
                               new Node(1.0/3.0)); // standard deviation == sqrt(1/9) = 1/3

    network->addNode(mu);

    // sigma node
    InverseGammaNode * variance = new InverseGammaNode("sigma",
                                                      0.25, 1.0/12.0,  // standard deviation = sqrt(1/144) = 1/12
                                                      new Node(11.0),
                                                      new Node(2.5));

    network->addNode(variance);


    // observations
    for(int i = 0; i < numPoints; i++)
    {        
        network->addNode(new NormalNode(QString("obs") + QString::number(i), data[i], 0.0, mu, variance, true));
    }

    return network;
}





/******************************************************************************************************************************************************
 *
 *   Burglar Alarm
 *
 *   The Burglar Alarm example is described on pages 493 and 494 in Russell and Norvig (Artificial Intelligence). If you don't have the book,
 *   look at slides 5 and 6 on this pdf. This network involves only DISCRETE NODES.
 *
 ******************************************************************************************************************************************************/

Network * NetworkFactory::buildBurglarNetwork()
{
    Network * network = new Network();

    // create nodes
    BernoulliDirectNode * burglary   = new BernoulliDirectNode("Burglary",   0.0, 1.0, false);
    BernoulliDirectNode * earthQuake = new BernoulliDirectNode("EarthQuake", 0.0, 1.0, false);
    BernoulliDirectNode * alarm      = new BernoulliDirectNode("Alarm",      0.0, 1.0, false);
    BernoulliDirectNode * johnCalls  = new BernoulliDirectNode("JohnCalls",  0.0, 1.0, false);
    BernoulliDirectNode * maryCalls  = new BernoulliDirectNode("MaryCalls",  1.0, 1.0, true);

    // build relationships
    burglary->setStateDistribution(1, 0.001); // T
    burglary->setStateDistribution(0, 0.999); // F

    earthQuake->setStateDistribution(1, 0.002); // T
    earthQuake->setStateDistribution(0, 0.998); // F

    alarm->addParentNode(burglary);
    alarm->addParentNode(earthQuake);      // BE
    alarm->setStateDistribution(3, 0.95);  // TT
    alarm->setStateDistribution(2, 0.94);  // TF
    alarm->setStateDistribution(1, 0.29);  // FT
    alarm->setStateDistribution(0, 0.001); // FF


    johnCalls->addParentNode(alarm);          // A
    johnCalls->setStateDistribution(1, 0.9);  // T
    johnCalls->setStateDistribution(0, 0.05); // F

    maryCalls->addParentNode(alarm);          // A
    maryCalls->setStateDistribution(1, 0.7);  // T
    maryCalls->setStateDistribution(0, 0.01); // F


    // add to network
    network->addNode(burglary);
    network->addNode(earthQuake);
    network->addNode(alarm);
    network->addNode(johnCalls);
    network->addNode(maryCalls);

    return network;
}



/******************************************************************************************************************************************************
 *
 *   Simple Bernoulli Network
 *
 *   The Burglar Alarm example is described on pages 493 and 494 in Russell and Norvig (Artificial Intelligence). If you don't have the book,
 *   look at slides 5 and 6 on this pdf. This network involves only DISCRETE NODES.
 *
 ******************************************************************************************************************************************************/

Network * NetworkFactory::buildSimpleBernoulli()
{
    Network * network = new Network();

    // create nodes
    BernoulliDirectNode * A  = new BernoulliDirectNode("A",  0.0, false);
    BernoulliDirectNode * B  = new BernoulliDirectNode("B",  0.0, false);
    BernoulliDirectNode * C  = new BernoulliDirectNode("C",  1.0, true);

    // build relationships
    A->setStateDistribution(1, 0.001); // T
    A->setStateDistribution(0, 0.999); // F

    B->addParentNode(A);
    B->setStateDistribution(1, 0.29);  // T
    B->setStateDistribution(0, 0.001); // F


    C->addParentNode(B);
    C->setStateDistribution(1, 0.9);  // T
    C->setStateDistribution(0, 0.05); // F

    // add to network
    network->addNode(A);
    network->addNode(B);
    network->addNode(C);

    return network;
}





/******************************************************************************************************************************************************
 *
 *   The file golfdataR.dat reports the scores at 42 professional golf tournaments involving 604 professional golfers. The Bayesian network for this
 *   example contains about 6000 nodes. With such a large network, it is difficult to tune mixing and burn parameters, and it takes a while to run
 *   simulations. The model is as follows:
 *
 *              An observation for Golfer i in Tournament j (obsnode[i][j]) is normally distributed with mean golfermean[i]+tournmean[j] and
 *              variance obsvar. The node tournmean[j] represents the difficulty (average score) of Tournament j, and the node golfermean[i]
 *              represents the skill of the golfer (how much better they score than the average). The prior distribution of the node obsvar
 *              is an Inverse Gamma with shape=83 and scale=1/.0014. Each golfermean[i] is Normal with mean 0 and variance hypergolfervar.
 *              Hypergolfervar is Inverse Gamma with shape=18 and scale=1/.015. Each tournmean[j] is Normal with mean=hypertournmean and
 *              var=hypertournvar. Hypertournmean is Normal with mean=72 and var=2 and hypertournvar is Inverse Gamma with shape=18 and
 *              scale=1/.015.
 *
 *   Note that the model and priors were created by a statistician who is very familiar with professional golf.
 *   The nodes of interest are the golfermeans. It is insightful to compare the posterior distributions for different golfers. We took the 5th, 50th,
 *   and 95th percentiles of samples in the posterior distribution for each golfer and printed these out in ranked order:
 *
 *   1: VijaySingh     -3.840661; 90% interval: (-4.354664, -3.288281)
 *   2: TigerWoods     -3.737515; 90% interval: (-4.568503, -2.850548)
 *   3: ErnieEls       -3.430582; 90% interval: (-4.313688, -2.480963)
 *   4: PhilMickelson  -3.353783; 90% interval: (-3.961659, -2.734419)
 *   5: StewartCink    -3.044881; 90% interval: (-3.647329, -2.476737)
 *   ...
 *   600: TommyAaron    4.883014; 90% interval: ( 2.666637,  7.017471)
 *   601: DerekSanders  5.420959; 90% interval: ( 3.484806,  7.799434)
 *   602: BobLohr       6.026131; 90% interval: ( 3.881693,  8.108147)
 *   603: ArnoldPalmer  6.226227; 90% interval: ( 4.355256,  8.178954)
 *   604: TimTims       7.978016; 90% interval: ( 5.905149, 10.156050)
 *
 ******************************************************************************************************************************************************/

Network * NetworkFactory::buildGolfersNetwork()
{
    Network * network = new Network();
    int nodeCount = 4;

    typedef InverseGammaNode VarianceNode;



 // ---------------------- build constant nodes ----------------------------

      double estimateAverage       = 72.8;
      double hyperTournMeanSd      = 2.0;
      double hyperVarSd            = 1.0;// 0.27; //.00405
      double meanSd                = 1.0;
      double observationVarianceSd = 1.0; //0.1162;



    // tournament hyper mean
      NormalNode * hyperTournMean = new NormalNode( "Tournament Hyper Mean",
                                                    estimateAverage,
                                                    hyperTournMeanSd,
                                                    new Node(72.0),
                                                    new Node(2.0) );
      network->addNode(hyperTournMean);

    // tournament hyper variance
      VarianceNode * hyperTournVar  = new VarianceNode ( "Tournament Hyper Variance",
                                                                3.0,
                                                                hyperVarSd,
                                                                new Node(18.0),
                                                                new Node(1.0/0.015));
      network->addNode(hyperTournVar);



    // hyper golfer variance
      VarianceNode * hyperGolferVar = new VarianceNode ("Golfer Hyper Var",
                                                               3.5,
                                                               hyperVarSd,
                                                               new Node(18.0),
                                                               new Node(1.0/0.015));
      network->addNode(hyperGolferVar);

    // observation variance
      VarianceNode * obsVar = new VarianceNode ("Observation Var",
                                                       3.1,
                                                       observationVarianceSd,
                                                       new Node(83.0),
                                                       new Node(1.0/0.0014));

      network->addNode(obsVar);




// ---------------------- parse data and build tournament / golfer / observation nodes ----------------------------


    NormalNode * tournamentNode  = NULL;    
    NormalNode * observationNode = NULL;
    CombinationNode * cNode      = NULL;

    int currentTournament = -1;

    QMap<QString, Node *> golfers;


    // parse contents
       QFile file("../SMCMC/golfdataR.dat"); //"../SMCMC/golfdataR_small.dat"); //
       if (!file.open(QIODevice::ReadOnly | QIODevice::Text))
           return NULL;      

       QTextStream in(&file);
       while (!in.atEnd())
       {
               QString line = in.readLine();
                       line = line.simplified();
               QStringList chunks = line.split(" ", QString::SkipEmptyParts);
               // name = 0, score = 1, tournament = 2
               int tournament = atoi(chunks.at(2).toUtf8().data()) - 1;
               double score   = atof(chunks.at(1).toUtf8().data());
               QString name   = chunks.at(0);

             // ================ new tournament ================

               if(tournament != currentTournament)
               {
                   currentTournament = tournament;
                   tournamentNode = new NormalNode(QString("Tournament ") + QString::number(tournament),
                                                   estimateAverage,
                                                   meanSd,
                                                   hyperTournMean,
                                                   hyperTournVar);

                   network->addNode(tournamentNode);
                   nodeCount++;
               }



             // ================ new golfer ================

               if(!golfers.contains(name))
               {
                    golfers[name] = new NormalNode( name,
                                                 0.0,
                                                 meanSd,
                                                 new Node(0.0),
                                                 hyperGolferVar);
                    network->addNode(golfers[name]);
                    nodeCount++;
               }



            // ================ new observation ================

               // create combination node to handle special case: tournamentMeanNode[i] + golferMeanNode[name]
               cNode = new CombinationNode(QString("Combo") + QString::number(nodeCount), 0.0, 0.0, true);
               cNode->addNode(tournamentNode);
               cNode->addNode(golfers.value(name));

               observationNode = new NormalNode(QString("Observation") + QString::number(nodeCount),
                                                score,
                                                0.0,
                                                cNode,
                                                obsVar,
                                                true);
               network->addNode(observationNode);
               nodeCount++;
       }

    file.close();

    assert(nodeCount == network->numNodes());

    return network;
}





/******************************************************************************************************************************************************
 *
 *  A wacky network:
 *
 *  A = NormalNode(mean=20, var=1)
 *  E = BetaNode(alpha=1, beta=1)
 *  B = GammaNode(shape=A**pi, invscale=7)
 *  D = BetaNode(alpha=A, beta=E)
 *  C = BernoulliNode(D: True, (1-D): False)
 *  F = PoissonNode(rate=D)
 *  G = NormalNode(mean=E, var=F)
 *
 ******************************************************************************************************************************************************/
Network * NetworkFactory::buildWackyNetwork()
{
    Network * network = new Network();

  /*NormalNode          * A = new NormalNode          ( "A", 20.0,  1.0, new Node(20.0), new Node(1.0) );
    BetaNode            * E = new BetaNode            ( "E",  0.5,  1.0, new Node(1.0),  new Node(1.0) );
    GammaPiNode         * B = new GammaPiNode         ( "B",  1750, 0.40816, A,          new Node(7.0) );
    BetaNode            * D = new BetaNode            ( "D",  0.9523,  1.0, A,              E );
    BernoulliDirectNode * C = new BernoulliDirectNode ( "C",  0.0,  1.0, D );
    PoissonNode         * F = new PoissonNode         ( "F",  1.0,  1.0, D );
    NormalNode          * G = new NormalNode          ( "G",  0.5,  1.0, E,              F );*/

    NormalNode          * A = new NormalNode          ( "A",  20.0,       1.0,         new Node(20.0), new Node(1.0) );
    BetaNode            * E = new BetaNode            ( "E",  0.5,        1.0,     new Node(1.0),  new Node(1.0) );
    GammaPiNode         * B = new GammaPiNode         ( "B",  1746.65944, 0.40816,     A,              new Node(7.0) );
    BetaNode            * D = new BetaNode            ( "D",  0.5,     1.0, A,              E );
    BernoulliDirectNode * C = new BernoulliDirectNode ( "C",  0.0,     1.0,  D );
    PoissonNode         * F = new PoissonNode         ( "F",  1.0,     1.0,      D );
    NormalNode          * G = new NormalNode          ( "G",  0.5,        1.0,      E,              F );

    network->addNode(A);
    network->addNode(E);
    network->addNode(B);
    network->addNode(D);
    network->addNode(C);
    network->addNode(F);
    network->addNode(G);

    return network;
}


Network * NetworkFactory::buildOrbitNetwork()
{
    QFile file("../SMCMC/orbit_small.dat");
    if (!file.open(QIODevice::ReadOnly | QIODevice::Text))
        return NULL;

    QTextStream in(&file);
    QStringList chunks = in.readLine().simplified().split(" ", QString::SkipEmptyParts);

    double massOrbital = atof(chunks.at(0).toUtf8().data());
    double massOrbiter = atof(chunks.at(1).toUtf8().data());


    // ================ mass of orbital node ================
    Node * orbitalMass = new Node(massOrbital);

    // ================ mass of orbiter node ================
    Node * orbiterMass = new Node(massOrbiter);


    int cnt = 0;
    double x, y;
    CombinationNode * accelMean;

    while (!in.atEnd())
    {
        chunks = in.readLine().simplified().split(" ", QString::SkipEmptyParts);

        // name = 0, score = 1, tournament = 2
        x = atof(chunks.at(0).toUtf8().data());
        y = atof(chunks.at(1).toUtf8().data());

        if(cnt % 3 == 0)
        {
           // accelMean = new CombinationNode();
        }

       // accelMean->addNode(pos);

        cnt++;
    }
}


Network * NetworkFactory::buildPopcornNetwork()
{
    Network * network = new Network();

    BetaNode * moistureMean = new BetaNode("MoistureMean", 0.8, 0.02, new Node(6.0), new Node(1.5));
    network->addNode(moistureMean);

    BetaNode * leakMean        = new BetaNode("LeakMean", 0.07, 0.01, new Node(1.5), new Node(19.0), true);

    QFile file("../SMCMC/popcorn.dat");
    if (!file.open(QIODevice::ReadOnly | QIODevice::Text))
        return NULL;

    QTextStream in(&file);
    QStringList chunks;
    double popped = 0.0;
    int i = 0;
    BernoulliDirectNode * observedKernel;
    BernoulliDirectNode * moisture;
    BernoulliDirectNode * leak;

    while (!in.atEnd())
    {

        moisture   = new BernoulliDirectNode("HasMoisture" + QString::number(i), 1.0, 1.0, moistureMean);
        network->addNode(moisture);

        leak = new BernoulliDirectNode("HasLeak" + QString::number(i), 1.0, 1.0, leakMean, false);
        network->addNode(leak);


        chunks = in.readLine().simplified().split(" ", QString::SkipEmptyParts);
        popped = atof(chunks.at(0).toUtf8().data());

        observedKernel = new BernoulliDirectNode(
                QString("Kernel ") + QString::number(i),
                    popped,
                    true);

        observedKernel->addParentNode(moisture);
        observedKernel->addParentNode(leak);

        observedKernel->setStateDistribution(3, 0.05);  // TT
        observedKernel->setStateDistribution(2, 0.95);  // TF
        observedKernel->setStateDistribution(1, 0.001); // FT
        observedKernel->setStateDistribution(0, 0.05);   // FF

        network->addNode(observedKernel);

        i++;
    }

    file.close();

    return network;
}


Network * NetworkFactory::buildChainNetwork(int numChild)
{
    Network * network = new Network();

    NormalNode * top    = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(10.0), new Node(1.0));
    network->addNode(top);
    for(int i = 0; i < numChild; i++)
    {
        top    = new NormalNode(Node::generateId(), 1.0, 1.0, top, new Node(1.0));
        network->addNode(top);
    }

    return network;
}


Network * NetworkFactory::buildYNetwork(int numChild)
{
    Network * network = new Network();

    NormalNode * top_left  = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(10.0), new Node(1.0));
    NormalNode * top_right = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(10.0), new Node(1.0));
    network->addNode(top_left);
    network->addNode(top_right);

    NormalNode * bottomi    = new NormalNode(Node::generateId(), 1.0, 1.0, top_left, top_right);
    network->addNode(bottomi);

    for(int i = 0; i < numChild; i++)
    {
        bottomi = new NormalNode(Node::generateId(), 1.0, 1.0, bottomi, new Node(1.0));
        network->addNode(bottomi);
    }

    return network;
}


Network * NetworkFactory::buildPeakNetwork()
{
    Network * network = new Network();

    NormalNode * top          = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(1.0), new Node(1.0));
    NormalNode * bottom_left  = new NormalNode(Node::generateId(), 1.0, 1.0, top, new Node(1.0));
    NormalNode * bottom_right = new NormalNode(Node::generateId(), 1.0, 1.0, top, new Node(1.0));

    network->addNode(top);
    network->addNode(bottom_left);
    network->addNode(bottom_right);

    return network;
}

Network * NetworkFactory::buildFunkySimple()
{
    Network * network = new Network();

    NormalNode * top_left     = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(1.0), new Node(1.0));
    NormalNode * top_right    = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(1.0), new Node(1.0));
    NormalNode * bottom_left  = new NormalNode(Node::generateId(), 1.0, 1.0, top_left, top_right);
    NormalNode * bottom_right = new NormalNode(Node::generateId(), 1.0, 1.0, top_left, top_right);

    network->addNode(top_left);
    network->addNode(top_right);
    network->addNode(bottom_left);
    network->addNode(bottom_right);

    return network;
}

Network * NetworkFactory::buildRandomNetwork(int numNodes, int numCrossLinks)
{
    Network * network = new Network();
    buildRandomNetwork(numNodes, network);

    QVector<Node*> * nodes = network->getNodes();
    int nodeCount = nodes->count();
    int indexFrom = 0, indexTo = 0;
    numCrossLinks = nodeCount < numCrossLinks ? nodeCount - 1 : numCrossLinks;
    for(int i = 0; i < numCrossLinks; i++)
    {
        while(indexFrom == indexTo)
        {
            indexFrom = rand() % nodeCount;
            indexTo = rand() % nodeCount;
        }
        nodes->at(indexFrom)->addParentNode(nodes->at(indexTo));
    }

    return network;
}

Node * NetworkFactory::buildRandomNetwork(int cnt, Network * network)
{
    Node * n;
    if(cnt == 0)
    {
        n = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(1.0), new Node(1.0));
        network->addNode(n);
        return n;
    }

    int state = rand() % 5;
    switch(state)
    {
        // add 2 parents
        case 0:
            n = new NormalNode(Node::generateId(), 1.0, 1.0, buildRandomNetwork(cnt - 1, network), buildRandomNetwork(cnt - 1, network));
            break;

        // add 1 parent left
        case 1:
        case 2:
            n = new NormalNode(Node::generateId(), 1.0, 1.0, buildRandomNetwork(cnt - 1, network), new Node(1.0));
            break;

        // add 1 parent right
        case 3:
        case 4:
            n = new NormalNode(Node::generateId(), 1.0, 1.0, new Node(1.0), buildRandomNetwork(cnt - 1, network));
            break;
    }

    int addChild = rand() % 50;

    network->addNode(n);

    if(addChild == 1)
    {
        Node * child;
        if(rand() % 2 == 0)
        {
            child = new NormalNode(Node::generateId(), 1.0, 1.0, n, buildRandomNetwork(cnt, network));
        } else {
            child = new NormalNode(Node::generateId(), 1.0, 1.0, buildRandomNetwork(cnt, network), n);
        }
        network->addNode(child);
    }

    return n;
}
